use syn::{File, Item, ItemUse, UseTree, Path};
use std::{fs, collections::HashSet, env};

/// 需要自动导入的常见库符号
const COMMON_CRATES: &[&str] = &[
    "use libc::strcpy;",
    "use libc::strdup;",
    "use libc::strlen;",
    "use libc::free;",
    "use libc::malloc;",
    "use std::assert;",
    "use std::cmp::Ordering;",
    "use std::cmp::min;",
    "use std::collections::HashMap;",
    "use std::collections::HashSet;",
    "use std::f64;",
    "use std::ffi::CStr;",
    "use std::ffi::CString;",
    "use std::fmt::Debug;",
    "use std::fs::File;",
    "use std::io;",
    "use std::io::stdin;",
    "use std::io::stdout;",
    "use std::io::Read;",
    "use std::io::Write;",
    "use std::ops::Add;",
    "use std::os::raw::c_char;",
    "use std::os::raw::c_void;",
    "use std::panic::AssertUnwindSafe;",
    "use std::path::Path;",
    "use std::process;",
    "use std::process::Command;",
    "use std::ptr;",
    "use std::mem;",
    "use std::ptr::null_mut;",
    "use std::rc::Rc;",
    "use std::str;",
    "use std::str::FromStr;",
    "use std::thread;",
    "use std::time::Duration;",
    "use std::env;",
];

/// 提取 `use` 语句中的路径，递归处理树结构
fn extract_use_paths(use_tree: &UseTree) -> Vec<String> {
    let mut paths = Vec::new();

    match use_tree {
        UseTree::Path(path) => {
            paths.push(path.ident.to_string());
        }
        UseTree::Name(name) => {
            paths.push(name.ident.to_string());
        }
        UseTree::Glob(_) => {} // 忽略 glob 使用
        UseTree::Group(group) => {
            for item in &group.items {
                paths.extend(extract_use_paths(item));
            }
        }
        UseTree::Rename(_) => {}
    }

    paths
}

/// 扫描文件内容，提取所有使用的符号
fn extract_used_symbols(file_path: &str) -> std::io::Result<HashSet<String>> {
    let content = fs::read_to_string(file_path)?;

    // 尝试解析 Rust 文件
    let syntax: Result<File, _> = syn::parse_str(&content);
    
    // 如果无法解析，则返回空值
    if syntax.is_err() {
        return Ok(HashSet::new());
    }

    let syntax = syntax.unwrap();
    let mut used_symbols = HashSet::new();

    for item in syntax.items {
        match item {
            Item::Use(item_use) => {
                for path in extract_use_paths(&item_use.tree) {
                    used_symbols.insert(path);
                }
            }
            Item::Fn(item_fn) => {
                // 提取函数签名中的标识符（例如函数调用的符号）
                let fn_name = item_fn.sig.ident.to_string();
                used_symbols.insert(fn_name);
            }
            Item::Struct(item_struct) => {
                // 提取结构体中的字段名称
                used_symbols.insert(item_struct.ident.to_string());
            }
            Item::Enum(item_enum) => {
                // 提取枚举中的变体名称
                used_symbols.insert(item_enum.ident.to_string());
            }
            _ => {}
        }
    }

    Ok(used_symbols)
}

/// 检查文件并添加缺少的 `use` 语句
fn add_missing_uses(file_path: &str) -> std::io::Result<()> {
    // 获取文件中的所有使用符号
    let used_symbols = extract_used_symbols(file_path)?;

    // 读取文件内容
    let mut content = fs::read_to_string(file_path)?;

    // 获取已存在的 use 语句
    let syntax: Result<File, _> = syn::parse_str(&content);
    if syntax.is_err() {
        return Ok(());
    }
    
    let syntax = syntax.unwrap();
    let mut existing_uses = HashSet::new();
    for item in syntax.items {
        if let Item::Use(item_use) = item {
            for path in extract_use_paths(&item_use.tree) {
                existing_uses.insert(path);
            }
        }
    }

    // 检查哪些符号需要导入，并将它们添加到 use 语句中
    let mut missing_uses = Vec::new();
    for symbol in &used_symbols {
        if !existing_uses.contains(symbol) { 
            // 符号未导入，检查是否在常见库中
            if COMMON_CRATES.contains(&symbol.as_str()) {
                missing_uses.push(symbol);
            }
        }
    }

    // 如果需要添加新的 `use` 语句，插入到文件顶部
    if !missing_uses.is_empty() {
        let mut new_content = String::new();
        for use_stmt in missing_uses {
            new_content.push_str(&format!("use {};\n", use_stmt));
        }
        new_content.push_str(&content);
        
        // 写入文件
        fs::write(file_path, new_content)?;
    }

    Ok(())
}

fn main() {
    // 从命令行获取文件路径
    let args: Vec<String> = env::args().collect();
    
    if args.len() != 2 {
        eprintln!("Usage: {} <file-path>", args[0]);
        return;
    }
    
    let file_path = &args[1];
    
    // 调用函数扫描并添加缺失的 use 语句
    if let Err(e) = add_missing_uses(file_path) {
        eprintln!("Error: {}", e);
    } else {
        println!("Successfully updated the file with missing use statements.");
    }
}

